# encoding:utf-8
# -----------------------------------------------------------
# "ML-Stealer: Stealing Prediction Functionality of Machine Learning Models with Mere Black-Box Access"
# @author: Shijie Wang, 2019.
# ------------------------------------------------------------

"""
Train victim model locally
"""
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier
from sklearn import preprocessing
from sklearn.calibration import CalibratedClassifierCV
from sklearn.externals import joblib
import pandas as pd
import numpy as np


def data_preprocessing(X_train_file, Y_train_file):
    X = preprocessing.scale(np.array(pd.read_csv(X_train_file, sep=',')))
    Y = np.array(pd.read_csv(Y_train_file, sep=','))
    return X, Y


def svm(X_train, Y_train):
    model = SVC(max_iter=10000)
    clf = CalibratedClassifierCV(model)
    clf.fit(X_train, Y_train)
    return clf


def randomforest(X_train, Y_train):
    model = RandomForestClassifier(n_estimators=100)
    clf = CalibratedClassifierCV(model)
    clf.fit(X_train, Y_train)
    return clf


def xgboost(X_train, Y_train):
    model = XGBClassifier(n_estimators=100)
    clf = CalibratedClassifierCV(model)
    clf.fit(X_train, Y_train)
    return clf


if __name__ == '__main__':
    X_train, Y_train = data_preprocessing('X_adult_train_1.csv', 'Y_adult_train_1.csv')
    svm_model = svm(X_train, Y_train)
    joblib.dump(svm_model, './victim_models/svm_part_adult.m')
    RF_model = randomforest(X_train, Y_train)
    joblib.dump(RF_model, './victim_models/rf_part_adult.m')
    XGboost_model = xgboost(X_train, Y_train)
    joblib.dump(XGboost_model, './victim_models/xgboost_part_adult.m')
